package com.rice.riceaiagent.memory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.jetbrains.annotations.NotNull;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.*;
import java.util.ArrayList;
import java.util.List;

/**
 * 自定义文件记忆存储
 * @author ricejson
 */
public class MyFileChatMemory implements ChatMemory {

    private static final Kryo kryo = new Kryo();
    private final File BASE_FILE;

    static {
        kryo.setRegistrationRequired(false);
        // 设置实例化策略
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public MyFileChatMemory(String dir) {
        this.BASE_FILE = new File(dir);
        // 如果根路径不存在就创建
        if (!BASE_FILE.exists()) {
            BASE_FILE.mkdirs();
        }
    }

    @Override
    public void add(String conversationId, Message message) {
        // 找到对应文件
        File file = getFile(conversationId);
        // 获取文件中原先的消息列表
        List<Message> messageList = getMessageList(file);
        // 追加新文件后写入
        messageList.add(message);
        saveMessageList(file, messageList);
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        // 获取目标文件
        File file = getFile(conversationId);
        // 获取原先的消息列表
        List<Message> messageList = getMessageList(file);
        // 追加新的消息
        messageList.addAll(messages);
        // 写入
        saveMessageList(file, messageList);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 根据conversationId 查找对应文件
        File file = getFile(conversationId);
        List<Message> messageList = getMessageList(file);
        return messageList.stream().skip(Math.max(0, messageList.size() - lastN)).toList();
    }

    @Override
    public void clear(String conversationId) {
        // 获取文件
        File file = getFile(conversationId);
        // 清除文件
        if (file.exists()) {
            file.delete();
        }
    }

    private void saveMessageList(File file, List<Message> messageList) {
        try(Output output = new Output(new FileOutputStream(file))) {
            kryo.writeObject(output, messageList);
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        }
    }

    private static List<Message> getMessageList(File file) {
        List<Message> messageList = new ArrayList<>();
        if (file.exists()) {
            // 读取文件
            try (Input input = new Input(new FileInputStream(file))) {
                messageList = kryo.readObject(input, ArrayList.class);
            } catch (FileNotFoundException e) {
                e.printStackTrace();
            }
        }
        return messageList;
    }

    @NotNull
    private File getFile(String conversationId) {
        File file = new File(this.BASE_FILE, conversationId + ".kryo");
        return file;
    }
}
